# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import logging
import requests
from concurrent.futures import ThreadPoolExecutor
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.uniswap_v2.kr_rest.product import UniswapProduct
from coin.exchange.uniswap_base.constants import (
    balance_checker_deployment_address,
    PRODUCTION_RPC_URL,
)
from coin.exchange.uniswap_v2.kr_rest.constants import (
    deployment_address,
)
from coin.exchange.uniswap_base.constants import (
    MAX_UINT_128,
    MAX_UINT_256,
    TOKEN_TO_ADDRESS,
    TOKEN_TO_DECIMAL_MULTIPLIER,
)
from coin.exchange.uniswap_v2.kr_rest.native_private_client_utils import (
    _load_contract,
)
import time
from web3 import Web3, HTTPProvider
from eth_utils.address import is_checksum_address
from enum import Enum


class UniswapOrderSide(Enum):
    BUY = "buy"
    SELL = "sell"


class UniswapOrderType(Enum):
    LIMIT = "limit"
    MARKET = "market"


class UniswapOrderStatus(Enum):
    LIVE = "live"
    FILLED = "filled"
    CANCELED = "cancelled"


def to_uniswap_order_side(order_side: OrderSide):
    assert isinstance(order_side, OrderSide), order_side
    return getattr(UniswapOrderSide, order_side.name)


def from_uniswap_order_side(order_side: UniswapOrderSide):
    assert isinstance(order_side, UniswapOrderSide), order_side
    return getattr(OrderSide, order_side.name)


def to_uniswap_order_type(order_type: OrderType):
    assert isinstance(order_type, OrderType), order_type
    return getattr(UniswapOrderSide, order_type.name)


def from_uniswap_order_type(order_type: UniswapOrderType):
    assert isinstance(order_type, UniswapOrderType), order_type
    return getattr(OrderType, order_type.name)


class UniswapV2NativePrivateClient(object):
    def __init__(
        self,
        wallet_address=None,
        uniswap_instance=None,
        acct=None,
        currency_symbol_list=None,
        *args,
        **kwargs,
    ):
        assert wallet_address is not None
        assert uniswap_instance is not None
        assert is_checksum_address(wallet_address), wallet_address

        self.currencies = currency_symbol_list
        if "ETH" in self.currencies or "WETH" in self.currencies:
            self.currencies += ["WETH", "ETH"]

        if "BTC" in self.currencies or "WBTC" in self.currencies:
            self.currencies += ["WBTC", "BTC"]

        self.currencies = list(set(self.currencies))
        self.token_list = [TOKEN_TO_ADDRESS['ethereum'][symb] for symb in self.currencies]
        self._executor = ThreadPoolExecutor(max_workers=20)

        self.wallet_address = wallet_address
        self.acct = acct

        self.uniswap_instance = (
            Web3(HTTPProvider(PRODUCTION_RPC_URL['ethereum']))
            if (uniswap_instance is None)
            else uniswap_instance[0]
        )  # for passing a reference

        self.current_nonce = self.uniswap_instance.eth.get_transaction_count(
            self.wallet_address
        )

        #self.approve_SwapRouter(0)
        #self.approve_SwapRouter()
        #assert False

    def _router(self):
        return _load_contract(
            self.uniswap_instance,
            "router02",
            deployment_address["UniswapV2Router02"],
        )

    router = property(fget=_router, fset=None, fdel=None, doc=None)

    def _addr_balance_checker(self):
        return _load_contract(
            self.uniswap_instance,
            "balanceChecker",
            balance_checker_deployment_address["mainnet"],
        )

    addr_balance_checker = property(
        fget=_addr_balance_checker, fset=None, fdel=None, doc=None
    )

    def swap_out(
        self,
        path,
        amount_out,
        amount_in_maximum=((1 << 256) - 1),
        deadline=None,
    ):
        if deadline is None:
            deadline = int(time.time()) + 2 * 60

        assert len(path) >= 2
        args = None
        tx_args = {
                "nonce": self.current_nonce,
                "from": self.wallet_address,
                "type": "0x2",
                "maxFeePerGas": self.uniswap_instance.eth.max_priority_fee
                + (2 * self.uniswap_instance.eth.get_block("latest")["baseFeePerGas"]),
                "maxPriorityFeePerGas": self.uniswap_instance.eth.max_priority_fee,
            }
        
        exec_func = self.router.functions.swapTokensForExactTokens

        if path[-1] == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            path[-1] = TOKEN_TO_ADDRESS['ethereum']["WETH"]
            exec_func = self.router.functions.swapTokensForExactETH
        
        if path[0] == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            path[0] = TOKEN_TO_ADDRESS['ethereum']["WETH"]
            ETH_USED = amount_in_maximum
            exec_func = self.router.functions.swapETHForExactTokens
            tx_args |= {
                "value" : ETH_USED
            }
            args = (amount_out, path, self.wallet_address, deadline)

        if args is None:
            args = (amount_out, amount_in_maximum, path, self.wallet_address, deadline)
        
        transaction = exec_func(
            *args
        ).build_transaction(tx_args)
        self.current_nonce += 1
        transaction["gas"] = int(
            self.uniswap_instance.eth.estimate_gas(transaction) * 1.2
        )
        signed_txn = self.uniswap_instance.eth.account.sign_transaction(
            transaction, private_key=self.acct.key
        )
        transaction = self.uniswap_instance.eth.send_raw_transaction(
            signed_txn.rawTransaction
        )
        receipt = self.uniswap_instance.eth.wait_for_transaction_receipt(transaction)
        return receipt

    def swap_in(
        self,
        path,
        amount_in,
        amount_out_minimum=0,
        deadline=None,
    ):
        if deadline is None:
            deadline = int(time.time()) + 2 * 60

        assert len(path) >= 2
        args = None
        tx_args = {
                "nonce": self.current_nonce,
                "from": self.wallet_address,
                "type": "0x2",
                "maxFeePerGas": self.uniswap_instance.eth.max_priority_fee
                + (2 * self.uniswap_instance.eth.get_block("latest")["baseFeePerGas"]),
                "maxPriorityFeePerGas": self.uniswap_instance.eth.max_priority_fee,
            }

        exec_func = self.router.functions.swapExactTokensForTokensSupportingFeeOnTransferTokens
        if path[-1] == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            path[-1] = TOKEN_TO_ADDRESS['ethereum']["WETH"]
            exec_func = self.router.functions.swapExactTokensForETHSupportingFeeOnTransferTokens
        
        if path[0] == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            path[0] = TOKEN_TO_ADDRESS['ethereum']["WETH"]
            ETH_USED = amount_in
            exec_func = self.router.functions.swapExactETHForTokensSupportingFeeOnTransferTokens
            tx_args |= {
                "value" : ETH_USED
            }
            args = (amount_in, path, self.wallet_address, deadline)

        if args is None:
            args = (amount_in, amount_out_minimum, path, self.wallet_address, deadline)

        transaction = exec_func(
            *args            
        ).build_transaction(
            tx_args
        )
        self.current_nonce += 1
        transaction["gas"] = int(
            self.uniswap_instance.eth.estimate_gas(transaction) * 1.2
        )
        signed_txn = self.uniswap_instance.eth.account.sign_transaction(
            transaction, private_key=self.acct.key
        )
        transaction = self.uniswap_instance.eth.send_raw_transaction(
            signed_txn.rawTransaction
        )
        receipt = self.uniswap_instance.eth.wait_for_transaction_receipt(transaction)
        return receipt


    def place_order_and_fill_or_fail(
        self,
        *,
        product: UniswapProduct,
        path = None,
        price: [float, str],
        quantity: [float, str],
        order_side: UniswapOrderSide,
        order_type: UniswapOrderType = UniswapOrderType.LIMIT,
        slippage: float = 0.01,              # Totally ignore slippage
    ):
        # a valid swap path should at least have on swapping pool/pair
        # uniswap_product -> quote and base
        
        # CCX special treatment for unknown high slippage problem
        
        trading_in_ccx = product.quote.native_currency == "CCX" or product.base.native_currency == "CCX"
        if trading_in_ccx:
            slippage = 0.2
        
        quote_address, base_address = (
            TOKEN_TO_ADDRESS['ethereum'][product.quote.native_currency],
            TOKEN_TO_ADDRESS['ethereum'][product.base.native_currency]
        )
            
        if quote_address == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            quote_address = TOKEN_TO_ADDRESS['ethereum']["WETH"]
            
        if base_address == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            base_address = TOKEN_TO_ADDRESS['ethereum']["WETH"]
        
        if path is None:
            path = [quote_address, base_address]
        assert len(path) >= 2
        assert path[0] != path[-1]
        
        if path is not None:
            if (quote_address != TOKEN_TO_ADDRESS['ethereum']["WETH"]) and (base_address != TOKEN_TO_ADDRESS['ethereum']["WETH"]):
                path = [quote_address, TOKEN_TO_ADDRESS['ethereum']["WETH"], base_address]
        # price = p_base / p_quote
        price = float(price)
        quantity = float(quantity)
        # exchange "1" quote <-> "price" base

        # BUY short quote long base
        # SELL  short base long quote
    
        # PATH = [CCX ETH USDT] equivalent [USDT ETH CCX]
        if order_side == UniswapOrderSide.BUY:
            # reform swap path
            # short quote and long base
            if path[0] != quote_address:
                path = path[::-1] # swap in quote and swap out base
            assert path[0] == quote_address
            assert path[-1] == base_address
            amount_out = int(quantity / TOKEN_TO_DECIMAL_MULTIPLIER[product.base.native_currency])
            amount_in_maximum = int(quantity * price * (1 + slippage) / TOKEN_TO_DECIMAL_MULTIPLIER[product.quote.native_currency])
            fut = self._executor.submit(
                self.swap_out,
                path,
                amount_out,
                MAX_UINT_128#, amount_in_maximum# MAX_UINT_256 if trading_in_ccx else amount_in_maximum
            )
            return fut

        if order_side == UniswapOrderSide.SELL:
            # reform swap path
            # short base and long quote
            if path[0] != base_address:
                path = path[::-1] # swap in base and swap out quote
            assert path[0] == base_address
            assert path[-1] == quote_address
            amount_in = int(quantity / TOKEN_TO_DECIMAL_MULTIPLIER[product.base.native_currency])
            amount_out_minimum = int(quantity * price * (1 - slippage) / TOKEN_TO_DECIMAL_MULTIPLIER[product.quote.native_currency])
            fut = self._executor.submit(
                self.swap_in,
                path,
                amount_in,
                0 #,amount_out_minimum# 0 if trading_in_ccx else amount_out_minimum
            )
            return fut

        raise NotImplementedError()


    def get_order(self, order_id: int):
        raise NotImplementedError()

    def cancel_order(self, order_id):
        raise NotImplementedError()

    def approve_SwapRouter(self, amount=MAX_UINT_256):
        logging.info("approve SwapRouter")
        contracts = [
            _load_contract(self.uniswap_instance, "erc20", Web3.to_checksum_address(addr))
            for addr in self.token_list if addr != TOKEN_TO_ADDRESS['ethereum']["USDT"]
        ]

        if TOKEN_TO_ADDRESS['ethereum']["USDT"] in self.token_list:
            contracts += [
                _load_contract(
                    self.uniswap_instance, 
                    "usdt", 
                    Web3.to_checksum_address(TOKEN_TO_ADDRESS['ethereum']["USDT"])
                )
            ]

        for instance in contracts:
            tx = instance.functions.approve(
                deployment_address["UniswapV2Router02"], amount
            ).build_transaction(
                {
                    "nonce": self.current_nonce,
                    "from": self.wallet_address,
                    "type": "0x2",
                    "maxFeePerGas": self.uniswap_instance.eth.max_priority_fee
                    + (
                        2
                        * self.uniswap_instance.eth.get_block("latest")["baseFeePerGas"]
                    ),
                    "maxPriorityFeePerGas": self.uniswap_instance.eth.max_priority_fee,
                }
            )
            self.current_nonce += 1
            tx["gas"] = int(self.uniswap_instance.eth.estimate_gas(tx) * 1.2)
            logging.info(tx)
            signed_txn = self.uniswap_instance.eth.account.sign_transaction(
                tx, private_key=self.acct.key
            )
            transaction = self.uniswap_instance.eth.send_raw_transaction(
                signed_txn.rawTransaction
            )
            recipt = self.uniswap_instance.eth.wait_for_transaction_receipt(transaction)


    def get_balances(self):
        logging.debug("Start get balance")
        req_timestamp = time.time()

        token_balance = {
            symbol: {"currency": symbol, "hold": 0, "volume": 0}
            for symbol in TOKEN_TO_ADDRESS['ethereum'].keys()
        }

        res_timestamp = time.time()

        address_coin_balances = self.addr_balance_checker.functions.balances(
            [self.wallet_address], self.token_list
        ).call()

        for token, balance in zip(list(self.currencies), address_coin_balances):
            token_balance[token]["volume"] = balance * TOKEN_TO_DECIMAL_MULTIPLIER.get(
                token, 1e-18
            )

        logging.debug(f"{token_balance.values()}")

        return RestUpdate(
            req_timestamp,
            res_timestamp,
            list(token_balance.values()),
            requests.codes.ok,
        )
